# -*- coding: utf-8 -*-
"""
Created on Wed May 26 2021
@author: Qingtian
@contributor: 59567

需要pip install shap以及ipywidgets(用于显示进度条)。由于shap的依赖包slicer与我们的slicer重名，需要修改我们的slicer为其他名字。
shapley_out函数为入口函数
"""

import shap
import joblib
import torch
import pandas as pd
from pandas import DataFrame, date_range
import numpy as np
from constructor.rnns_pytorch import RNNs
from constructor.tcn import TCN
from torch import tensor
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']  # 显示中文
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号


def load_model(model_path):
    exported_model = torch.load(model_path)
    if s_dict_global['model'] == 'RNN':
        model = RNNs(model_name=s_dict_global['model'], input_size=int(s_dict_global['n_features']),
                     hidden_size=int(s_dict_global['RNN_hidden_size']),
                     num_layers=int(s_dict_global['RNN_num_layers']), nonlinearity=s_dict_global['RNN_nonlinearity'],
                     bias=s_dict_global['RNN_bias'],
                     batch_first=s_dict_global['RNN_batch_first'], dropout=int(s_dict_global['RNN_dropout']),
                     bidirectional=s_dict_global['RNN_bidirectional'],
                     time_step=int(s_dict_global['x_time_step']), output_size=int(s_dict_global['y_time_step']))
        print('RNN模型加载完成')
    elif s_dict_global['model'] == 'GRU':
        model = RNNs(model_name=s_dict_global['model'], input_size=int(s_dict_global['n_features']),
                     hidden_size=s_dict_global['GRU_hidden_size'],
                     num_layers=s_dict_global['GRU_num_layers'], bias=s_dict_global['GRU_bias'],
                     batch_first=s_dict_global['GRU_batch_first'], dropout=s_dict_global['GRU_dropout'],
                     bidirectional=s_dict_global['GRU_bidirectional'],
                     time_step=s_dict_global['x_time_step'], output_size=s_dict_global['y_time_step'])
        print('GRU模型加载完成')
    elif s_dict_global['model'] == 'LSTM':
        model = RNNs(model_name=s_dict_global['model'], input_size=int(s_dict_global['n_features']),
                     hidden_size=s_dict_global['LSTM_hidden_size'],
                     num_layers=s_dict_global['LSTM_num_layers'], bias=s_dict_global['LSTM_bias'],
                     batch_first=s_dict_global['LSTM_batch_first'], dropout=s_dict_global['LSTM_dropout'],
                     bidirectional=s_dict_global['LSTM_bidirectional'],
                     time_step=s_dict_global['x_time_step'], output_size=s_dict_global['y_time_step'])
        print('LSTM模型加载完成')
    elif s_dict_global['model'] == 'TCN':
        model = TCN(input_size=int(s_dict_global['n_features']), output_size=1, num_channels=[30] * 8,
                    kernel_size=2, dropout=0)
        print('TCN模型加载完成')
    else:
        raise Exception('暂不支持该深度学习模型')
    model.load_state_dict(exported_model)
    return model


def model_predict(x):
    model = load_model(model_path_global)
    output = []
    x_array = np.array(x)
    for i in range(len(x_array)):
        x_test = x_array[i, :]
        x_test_reshape = x_test.reshape((1, 1, len(x_test)))
        x_test_tensor = tensor(x_test_reshape, dtype=torch.float32)
        out = model(x_test_tensor)
        out = out.detach().numpy()
        output.append(out[0][0])
    return np.array(output)


def shapley_out(x: DataFrame, y: DataFrame, summary: dict, model_path: str, explain_recent=False) -> DataFrame:
    """
    计算模型日度还原的shapley值。自模型生效月开始，至还原最新日结束。
    Calculate shap values of the model. From the first day of the model in effect to the last day of available input.

    Parameters:
      x - 含有日度还原最新一期的X值的Dataframe。A dataframe of X values, containing needed inputs for predictions.
      y - (未使用)y值的Dataframe。A dataframe of y values, containing targets. Currently unused.
      summary - 模型的summary字典。Dictionary containing information of the model.
      model_path - 模型的保存路径。Path of the exported model.
      explain_recent - 是否只解释最近7天。默认为自模型生效月开始，至还原最新日结束。

    Returns:
        shapley值的dateframe。A dateframe of shapley values. Features in columns, dates as indexes.

    Raises:
        KeyError - raises an exception
    """

    # model_path_global和s_dict_global为全局变量，其他两个函数使用
    global model_path_global
    global s_dict_global
    model_path_global = model_path
    s_dict = summary.iloc[-1, :].to_dict()
    s_dict_global = s_dict
    y = y.loc[:, [s_dict['y_name']]]
    x_features = summary.iloc[-1]['features']

    if explain_recent == False:
        # x_explain包含自模型生效月开始，至还原最新日结束的x
        end_datetime_month = pd.date_range(start=pd.to_datetime(s_dict['end_datetime']).replace(day=1),
                                           end=x.last_valid_index(), freq="D")
        x_explain = x.loc[end_datetime_month, x_features]
    else:
        # x_explain包含最近的7天
        end_datetime_month = pd.date_range(end=x.last_valid_index(), periods=7, freq="D")
        x_explain = x.loc[end_datetime_month, x_features]

    if model_path.split('.')[-1] == 'pth':
        # 深度学习模型使用自己写的model_predict函数
        explainer = shap.Explainer(model_predict, x_explain)
        print('解释器初始化完成')
    elif model_path.split('.')[-1] == 'joblib':
        # sklearn模型使用自带的model.predict函数
        exported_model = joblib.load(model_path)
        model = exported_model.model
        explainer = shap.Explainer(model.predict, x_explain)
        print('解释器初始化完成')
    else:
        print('模型文件后缀名错误')
    global shap_values
    shap_values = explainer(x_explain)
    shap.plots.beeswarm(shap_values, )
    shap.plots.force(shap_values[-1], text_rotation=-30, matplotlib=True)
    shap_frame = pd.DataFrame(shap_values.values, index=end_datetime_month, columns=shap_values.feature_names)
    return shap_frame


if __name__ == "__main__":
    # 输入修改,只用改这里
    y = pd.read_excel(r"D:\2021_05_08\input\y_plus_集合.xlsx", index_col=0)
    summary = pd.read_pickle(r"C:\Users\huang\Desktop\工业增加值_当月同比\summary.pkl")
    model_path = r"C:\Users\huang\Desktop\工业增加值_当月同比\model.pth"
    X_full = pd.read_pickle(r"D:\2021_05_18\input\Baidu_index_2021-05-25.pkl")

    shap_frame = shapley_out(x=X_full, y=y, summary=summary, model_path=model_path, explain_recent=True)
